import torch

from tqdm import tqdm
from typing import Dict, List, Tuple
from data_process.structure import Point3D, RadarPoint, RtkData


class CoordinateTransformer:
    EARTH_A = 6378137.0
    EARTH_E2 = 6.69437999014e-3
    DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    @staticmethod
    def _latlonh_to_ecef(lat: float, lon: float, h: float) -> torch.Tensor:
        lat_rad = torch.deg2rad(torch.tensor(lat, device=CoordinateTransformer.DEVICE, dtype=torch.float64))
        lon_rad = torch.deg2rad(torch.tensor(lon, device=CoordinateTransformer.DEVICE, dtype=torch.float64))
        sin_lat = torch.sin(lat_rad)
        cos_lat = torch.cos(lat_rad)
        N = CoordinateTransformer.EARTH_A / torch.sqrt(1 - CoordinateTransformer.EARTH_E2 * sin_lat**2)
        return torch.stack([
            (N + h) * cos_lat * torch.cos(lon_rad),
            (N + h) * cos_lat * torch.sin(lon_rad),
            (N * (1 - CoordinateTransformer.EARTH_E2) + h) * sin_lat
        ])

    @staticmethod
    def _ecef_to_enu(ecef: torch.Tensor, ref_ecef: torch.Tensor, ref_lat: float, ref_lon: float) -> torch.Tensor:
        delta = ecef - ref_ecef
        lat_rad = torch.deg2rad(torch.tensor(ref_lat, device=CoordinateTransformer.DEVICE, dtype=torch.float64))
        lon_rad = torch.deg2rad(torch.tensor(ref_lon, device=CoordinateTransformer.DEVICE, dtype=torch.float64))
        row1 = torch.tensor([-torch.sin(lon_rad), torch.cos(lon_rad), 0], 
                        device=CoordinateTransformer.DEVICE, dtype=torch.float64)
        row2 = torch.tensor([
            -torch.sin(lat_rad)*torch.cos(lon_rad),
            -torch.sin(lat_rad)*torch.sin(lon_rad),
            torch.cos(lat_rad)
        ], device=CoordinateTransformer.DEVICE, dtype=torch.float64)
        row3 = torch.tensor([
            torch.cos(lat_rad)*torch.cos(lon_rad),
            torch.cos(lat_rad)*torch.sin(lon_rad),
            torch.sin(lat_rad)
        ], device=CoordinateTransformer.DEVICE, dtype=torch.float64)
        
        R = torch.stack([row1, row2, row3])
        return torch.matmul(R, delta)

    @staticmethod
    def _euler_to_rotation(roll: float, pitch: float, yaw: float) -> torch.Tensor:
        roll_rad = torch.deg2rad(torch.tensor(roll, device=CoordinateTransformer.DEVICE, dtype=torch.float64))
        pitch_rad = torch.deg2rad(torch.tensor(pitch, device=CoordinateTransformer.DEVICE, dtype=torch.float64))
        yaw_rad = torch.deg2rad(torch.tensor(yaw, device=CoordinateTransformer.DEVICE, dtype=torch.float64))
        
        Rx = torch.tensor([
            [1, 0, 0],
            [0, torch.cos(pitch_rad), -torch.sin(pitch_rad)],
            [0, torch.sin(pitch_rad), torch.cos(pitch_rad)]
        ], device=CoordinateTransformer.DEVICE, dtype=torch.float64)
        
        Ry = torch.tensor([
            [torch.cos(roll_rad), 0, torch.sin(roll_rad)],
            [0, 1, 0],
            [-torch.sin(roll_rad), 0, torch.cos(roll_rad)]
        ], device=CoordinateTransformer.DEVICE, dtype=torch.float64)
        
        Rz = torch.tensor([
            [torch.cos(yaw_rad), -torch.sin(yaw_rad), 0],
            [torch.sin(yaw_rad), torch.cos(yaw_rad), 0],
            [0, 0, 1]
        ], device=CoordinateTransformer.DEVICE, dtype=torch.float64)
        
        return torch.matmul(Ry, torch.matmul(Rx, Rz))

    @staticmethod
    def _get_transform_matrix(
        current_rtk: RtkData,
        ref_rtk: RtkData,
        ref_ecef: torch.Tensor,
        reverse: bool = False,
    ) -> torch.Tensor:
        current_ecef = CoordinateTransformer._latlonh_to_ecef(
            current_rtk.latitude, 
            current_rtk.longitude, 
            current_rtk.height
        )
        delta_enu = CoordinateTransformer._ecef_to_enu(
            current_ecef, ref_ecef, 
            ref_rtk.latitude, 
            ref_rtk.longitude
        )
        R_ref = CoordinateTransformer._euler_to_rotation(ref_rtk.roll, ref_rtk.pitch, ref_rtk.heading)
        R_cur = CoordinateTransformer._euler_to_rotation(current_rtk.roll, current_rtk.pitch, current_rtk.heading)

        R = R_ref.T @ R_cur

        delta_enu = torch.matmul(R_ref.T, delta_enu)

        if reverse:
            R = R.T
            delta_enu = -torch.matmul(R, delta_enu)
            
        T = torch.eye(4, device=CoordinateTransformer.DEVICE, dtype=torch.float64)
        T[:3, :3] = R
        T[:3, 3] = delta_enu
        return T

    @staticmethod
    def transform_points(
        points_data: Dict[str, List[Point3D]], 
        rtk_data: Dict[str, RtkData],
        ref_rtk: RtkData,
        reverse: bool = False,
    ) -> Tuple[Dict[str, List[Point3D]], RtkData]:
        ref_ecef = CoordinateTransformer._latlonh_to_ecef(
            ref_rtk.latitude, ref_rtk.longitude, ref_rtk.height
        ).to(CoordinateTransformer.DEVICE)
        
        results = {}
        for frame_id in tqdm(points_data.keys(), desc="Transforming coordinate"):
            if frame_id not in rtk_data:
                continue

            current_rtk = rtk_data[frame_id]
            T = CoordinateTransformer._get_transform_matrix(current_rtk, ref_rtk, ref_ecef, reverse)
            
            points = points_data[frame_id]
            if len(points) == 0:
                results[frame_id] = []
                continue
            coords = torch.tensor(
                [[p.x, p.y, p.z, 1] for p in points],
                device=CoordinateTransformer.DEVICE,
                dtype=torch.float64
            )
            transformed = (coords @ T.T).cpu().numpy()
            
            results[frame_id] = [
                p.__class__(
                    x=transformed[i][0],
                    y=transformed[i][1],
                    z=transformed[i][2],
                    **{k:v for k,v in p.__dict__.items() if k not in ('x','y','z')}
                ) for i, p in enumerate(points)
            ]
            
        return results
    
    @staticmethod
    def debug():
        ref_rtk = RtkData(
            timeStamp=0,
            latitude=39.9,
            longitude=116.4,
            height=50.0,
            roll=0, 
            pitch=0, 
            heading=0
        )

        current_rtk = RtkData(
            timeStamp=0,
            latitude=39.9,
            longitude=116.399883,  # 西 10m
            height=50.0,
            roll=0, 
            pitch=0, 
            heading=0
        )

        points_data = {
            "frame_0": [
                RadarPoint(0, 10, 0, strength=1.0, velocity=0.0, cohen_factor=0.5)
            ]
        }

        transformed_data = CoordinateTransformer.transform_points(
            points_data=points_data,
            rtk_data={"frame_0": current_rtk},
            ref_rtk=ref_rtk,
            reverse=False
        )

        transformed_point = transformed_data["frame_0"][0]
        print(f"new: ({transformed_point.x}, {transformed_point.y}, {transformed_point.z})")
